package com.ztom.v2;

import java.util.HashMap;
import java.util.Map;

/**
 * @author ZhangTao
 */
public class Code48PathSum3 {

    public class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;

        TreeNode() {
        }

        TreeNode(int val) {
            this.val = val;
        }

        TreeNode(int val, TreeNode left, TreeNode right) {
            this.val = val;
            this.left = left;
            this.right = right;
        }
    }

    public int pathSum(TreeNode root, long targetSum) {
        if (root == null) {
            return 0;
        }

        // 只能从父节点到子节点的路径
        return countCurNode(root, targetSum) + pathSum(root.left, targetSum) + pathSum(root.right, targetSum);
    }

   public int countCurNode(TreeNode node, long rest) {
       if (node == null) {
           return 0;
       }
       int count = 0;
       if (node.val == rest) {
           count++;
       }
       count += countCurNode(node.left, rest - node.val);
       count += countCurNode(node.right, rest - node.val);
       return count;
   }


    public int pathSum1(TreeNode root, long targetSum) {
        if (root == null) {
            return 0;
        }
        // 前缀和优化, 可能出现多次, value 记录的是次数
        Map<Long, Integer> preSumMap = new HashMap<>();
        preSumMap.put(0L, 1);
        return process(root, preSumMap, 0, targetSum);
    }

    public int process(TreeNode node, Map<Long, Integer> preSumMap, long curSum, long targetSum) {
        if (node == null) {
            return 0;
        }

        // 当前节点前缀和
        curSum += node.val;
        // 累加符合条件的前缀和 (转换成两数之和问题了)
        int count = preSumMap.getOrDefault(curSum - targetSum, 0);

        // 取当前节点
        preSumMap.put(curSum, preSumMap.getOrDefault(curSum, 0) + 1);
        // 递归左右子树
        count += process(node.left, preSumMap, curSum, targetSum);
        count += process(node.right, preSumMap, curSum, targetSum);
        // 回溯恢复现场: 只能父节点到子节点的路径
        preSumMap.put(curSum, preSumMap.get(curSum) - 1);
        return count;
    }
}
